#include "IRpropPlus.h"
#include "ReClaM/Rprop.h"
#include <limits>
#include "../clUtils.h"
#include "../BindKernelArgs.h"

using namespace std;

learn::IRpropPlus::IRpropPlus(cl_command_queue _queue, int modelParameterCount, float initialDeltaW) 
	: stepIncreaseFactor(1.2f), stepDecreaseFactor(0.5f), 
	stepMaxIncrement(numeric_limits<float>::max()), stepMinIncrement(0), 
	prevError(numeric_limits<float>::max()),
	queue(_queue)
{
	cl_context context = getContext(queue);
	cl_device_id deviceId = getDeviceId(queue);
		
	cl_program program = buildProgramFromFile(context, deviceId, "rprop.cl");		
	cl_int errorCode;
	
    iRpropPlus = clCreateKernel(program, "iRpropPlus", &errorCode);
	CL_CHECK(errorCode);
	
	delta = ClMem<float>(queue, modelParameterCount, initialDeltaW);
	deltaW = ClMem<float>(queue, modelParameterCount);
	prevdEdW = ClMem<float>(queue, modelParameterCount);
}

learn::IRpropPlus::~IRpropPlus() {	
}

float sign(float x) {
	if (x > 0) return 1; 
	if (x < 0) return -1; 
	return 0;
}	

float learn::IRpropPlus::getStepIncreaseFactor() {
	return stepIncreaseFactor;
}

void learn::IRpropPlus::setStepIncreaseFactor(float value) {
	stepIncreaseFactor = value;
}

float learn::IRpropPlus::getStepDecreaseFactor() {
	return stepDecreaseFactor;
}

void learn::IRpropPlus::setStepDecreaseFactor(float value) {
	stepDecreaseFactor = value;
}

float learn::IRpropPlus::getStepMaxIncrement() {
	return stepMaxIncrement;
}

void learn::IRpropPlus::setStepMaxIncrement(float value) {	
	stepMaxIncrement = value;
}

float learn::IRpropPlus::getStepMinIncrement() {
	return stepMinIncrement;
}

void learn::IRpropPlus::setStepMinIncrement(float value) {
	stepMinIncrement = value;
}

ClEvent learn::IRpropPlus::iRpropPlusAsync(ClMem<float> parameters, ClMem<float> dEdW, float currentError) {
		
	BindKernelArgs(iRpropPlus) 
			<< dEdW << prevdEdW 
			<< delta << deltaW
			<< parameters
			<< currentError << prevError
			<< stepDecreaseFactor << stepIncreaseFactor
			<< stepMinIncrement << stepMaxIncrement;

	size_t count = parameters.count();
    cl_event event;
	CL_CHECK(clEnqueueNDRangeKernel(queue,
			iRpropPlus, 1, NULL,
			&count, NULL,
			0, NULL, &event));
	return ClEvent(event);
}

float learn::IRpropPlus::optimize(Model &model, ErrorFunction &errorFunction, OptimizationSupervisor &data) {
	if (model.getInputCount() != data.getInputLength()) {
		throw ArgumentException("model input count should be equal to train data input count");
	}
	
	if (model.getOutputCount() != data.getTargetLength()) {
		throw ArgumentException("model output count should be equal to train data output count");
	}
		
	int parameterCount = model.getParameterCount();
	
	vector<float> dEdW(parameterCount);	
	float currentError = errorFunction.getErrorDerivative(model, data, &dEdW[0]);
	ClMem<float> devicedEdW(queue, parameterCount);
	
	vector<float> parameters(parameterCount);
	model.getParameters(&parameters[0], 0, parameterCount);
	ClMem<float> deviceParams(queue, parameterCount);
	
	WAIT_ALL(
			devicedEdW.beginWrite(&dEdW[0], 0, parameterCount),
			deviceParams.beginWrite(&parameters[0], 0, parameterCount));
	
	iRpropPlusAsync(deviceParams, devicedEdW, currentError).wait();	
	deviceParams.beginRead(&parameters[0], 0, parameterCount).wait();
	
	prevError = currentError;
	model.setParameters(&parameters[0], 0, parameterCount);
	return currentError;
}